"""
联邦学习服务
用于与现有系统集成
"""

from typing import List, Dict, Any, Callable
from src.research_core.federated_learning import LangChainFederatedServer
from src.research_core.federated_client import start_federated_client
from src.research_core.model_manager import ModelManager, ModelType
from src.research_core.model_evaluator import ModelEvaluator


class FederatedLearningService:
    """联邦学习服务"""
    
    def __init__(self):
        self.model_manager = ModelManager()
        self.model_evaluator = ModelEvaluator()
        self.servers = {}
        self.clients = {}
    
    def create_federated_server(self, model_type: ModelType, min_fit_clients: int = 2, min_available_clients: int = 2) -> LangChainFederatedServer:
        """创建联邦学习服务器"""
        server = LangChainFederatedServer(model_type, min_fit_clients, min_available_clients)
        self.servers[model_type.value] = server
        return server
    
    def start_federated_training(self, model_type: ModelType, server_address: str = "0.0.0.0:8080"):
        """启动联邦训练"""
        if model_type.value in self.servers:
            server = self.servers[model_type.value]
            server.start_server(server_address)
        else:
            raise ValueError(f"未找到模型类型 {model_type.value} 的服务器")
    
    def register_federated_client(self, model_type: ModelType, train_data, test_data, server_address: str = "localhost:8080"):
        """注册联邦学习客户端"""
        client_key = f"{model_type.value}_{server_address}"
        self.clients[client_key] = {
            "model_type": model_type,
            "train_data": train_data,
            "test_data": test_data,
            "server_address": server_address
        }
        return client_key
    
    def start_client(self, client_key: str):
        """启动客户端"""
        if client_key in self.clients:
            client_info = self.clients[client_key]
            start_federated_client(
                client_info["model_type"],
                client_info["train_data"],
                client_info["test_data"],
                client_info["server_address"]
            )
        else:
            raise ValueError(f"未找到客户端 {client_key}")


class FederatedLearningIntegration:
    """联邦学习与现有系统集成"""
    
    def __init__(self):
        self.federated_service = FederatedLearningService()
        self.model_manager = ModelManager()
        self.model_evaluator = ModelEvaluator()
    
    async def setup_federated_training(self, node_configs: List[Dict]):
        """设置联邦训练"""
        for config in node_configs:
            node_id = config["node_id"]
            model_type = config["model_type"]
            data_loader = config["data_loader"]
            
            # 获取模型实例
            model = self.model_manager.get_model(model_type)
            
            # 注册节点
            # 注意：这里需要根据实际数据加载器实现进行调整
            train_data = data_loader() if callable(data_loader) else data_loader
            test_data = config.get("test_data", train_data)
            
            self.federated_service.register_federated_client(
                model_type, train_data, test_data
            )
    
    async def run_model_optimization(self, model_type: ModelType, rounds: int = 10):
        """运行模型优化"""
        # 创建服务器
        server = self.federated_service.create_federated_server(model_type)
        
        # 启动训练（在实际应用中，这应该在单独的进程中运行）
        # server.start_server()
        
        # 这里返回模拟结果
        return {"status": "federated_training_started", "model_type": model_type.value, "rounds": rounds}
    
    def get_optimized_model(self, model_type: ModelType):
        """获取优化后的模型"""
        # 返回经过联邦学习优化的模型
        return self.model_manager.get_model(model_type)